import math
import torch
import torch.nn as nn

"""
SpectFormers

Resource code: https://badripatro.github.io/SpectFormers/

Paper: https://arxiv.org/pdf/2304.06446

Intro:
使用傅立叶变换实现的光谱层来捕捉架构初始层中的相关特征。此外，在网络的深层使用多头自注意力。
SpectFormer架构简单，它将图像标记转换到傅立叶域，
然后使用可学习的权重参数应用门控技术，最后进行傅立叶逆变换以获取信号。SpectFormer结合了光谱注意力和多头注意力。
"""
# 定义自定义的 Transpose 模块，用于在 nn.Sequential 中进行张量维度交换
class Transpose(nn.Module):
    def __init__(self, dim0, dim1):
        super().__init__()
        self.dim0 = dim0  # 第一个要交换的维度
        self.dim1 = dim1  # 第二个要交换的维度

    def forward(self, x):
        # 在前向传播中交换指定的维度
        return x.transpose(self.dim0, self.dim1)

# 定义频谱门控网络（Spectral Gating Network）
class SpectralGatingNetwork(nn.Module):
    def __init__(self, dim, h=8, w=5):
        super().__init__()
        # 初始化可学习的复权重参数，形状为 (h, w, dim, 2)
        # 最后一个维度大小为 2，用于表示实部和虚部
        self.complex_weight = nn.Parameter(
            torch.randn(h, w, dim, 2, dtype=torch.float32) * 0.02
        )
        self.h = h  # 频域中的高度维度
        self.w = w  # 频域中的宽度维度

    def forward(self, x, spatial_size=None):
        B, N, C = x.shape  # 获取输入的批次大小、令牌数量和特征维度

        if spatial_size is None:
            # 计算空间维度
            a = b = int(math.sqrt(N))
        else:
            a, b = spatial_size

        # 计算频域尺寸
        h = a
        w = b // 2 + 1

        # 将输入重塑为 (B, a, b, C)，以表示空间维度
        x = x.view(B, a, b, C)
        x = x.to(torch.float32)  # 确保数据类型为 float32，兼容 FFT 操作

        # 对 x 进行二维实数傅立叶变换，得到频域表示
        x = torch.fft.rfft2(x, dim=(1, 2), norm='ortho')

        # 将可学习的复权重参数转换为复数形式，形状为 (h, w, C)
        weight = torch.view_as_complex(self.complex_weight)
        weight = weight.unsqueeze(0)  # 扩展批次维度，形状变为 (1, h, w, C)

        # 在频域中逐元素相乘，实现频谱门控
        x = x * weight  # x 的形状为 (B, a, b_freq, C)

        # 对 x 进行二维逆傅立叶变换，转换回空间域
        x = torch.fft.irfft2(x, s=(a, b), dim=(1, 2), norm='ortho')

        # 将 x 重塑回原始形状 (B, N, C)
        x = x.reshape(B, N, C)

        return x

# 主程序，测试 SpectralGatingNetwork 模块
if __name__ == '__main__':
    # 自动检测设备
    DEVICE = 'mps' if torch.backends.mps.is_available() else 'cpu'  # 支持 MPS 则使用，否则使用 CPU
    print(f"Using device: {DEVICE}")
    block = SpectralGatingNetwork(64).to(DEVICE)
    input = torch.rand(1, 64, 64).to(DEVICE)
    output = block(input)
    print(f"Input size: {input.size()}, Output size: {output.size()}")